# -*- coding: utf-8 -*-
import time
import sys
import json
import threading
import socket
from math import ceil
import numpy as np
import os
os.environ['TF_CPP_MIN_LOG_LEVEL']='2'
import tensorflow as tf


class FLserver:
    #核心部分--控制
    command_id=set()
    isroot=False
    running=False
    th_listen=None #监听端口的线程
    state_lock = threading.Lock()
    state=0#工作状态 0未初始化 1正常工作 2正在训练/等待下层聚合 3向上层继续聚合/完成聚合 4接收、发送新模型、更新模型
    isdebug=False#调试模式下不连接mqtt
    flag_start_send=False#是否开始发送

    #核心部分--缓存
    local_ip=None
    data_bytes=None #发送的梯度数据
    send_bytes_once=None
    cloud_server=None #云服务器，向云服务器发送状态信息
    up_server=None
    local_model=None
    temp_model=None
    n_data_layer=None#有数据的层数，creat_model时初始化
    model_data_layer_index=None#有数据的层数，creat_model时初始化
    received_buffer_new=[]#接收新模型的缓存,是list（数据包发送总数）

    n_down_server=None
    down_server=None
    received_buffer=[]#接收模型的缓存,是list（下层服务器个数，数据包发送总数）
    is_buffer_initialized=[]#0未初始化不知道是否参与，1已经初始化，-1不参与
    received_weight=[]#接收各个ip下数据的权重，在发送最后一个数据时发送

    #参数部分
    wait_time=3    #等待时间,接收数据间隔超过这个时间则认为丢包
    send_speed=20   #发送数据的速度M/s，不能超过21.7
    aggregate_waite_time=300 #等待时间,开始聚合后这个时间下层ip还没有开始传输数据则放弃
    port=None
    local_weight=1#本地权重，如果不参与训练则为0
    isparticipant=False#是否参与本轮训练

    #应用部分
    client=None    #mqtt client
    mqtt_server=None
    x_train=[]#储存图片
    y_train=[]#储存标签

    #---------------主要函数、控制相关-----------------
    def listen(self):
        server_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        server_socket.bind(("", self.port)) 
        while(self.running):
            server_socket.settimeout(3)  #每3秒钟检测是否关闭
            try:
                receive_data, client = server_socket.recvfrom(1400)
                #print("来自客户端%s,发送的%s\n" % (client, receive_data))  #打印接收的内容
                josn_end=receive_data.index(b"}")+1
                receive_data2=receive_data[josn_end:]
                json_data = json.loads(receive_data[:josn_end].decode('utf-8'))
                if "feedback_id" in json_data:
                    if not self.feedback(client,json_data):
                        json_data={}#已经处理过的任务就跳过
                if "order" in json_data:
                    if (json_data['order']=="aggregation_data"):#下层上传的模型
                        self.receive_model(json_data,receive_data2)#接收下层上传的模型
                    elif (json_data['order']=="update_data"):#上层向下层传输的新模型
                        self.receive_update_model(json_data,receive_data2)#在第一次接收数据时调用startupdate,因为更新模型的数据传输是从上往下的，不需要确认命令
                    elif (json_data['order']=="aggregation_request"):#开始聚合的请求
                        if(self.isroot):
                            th_start_aggregate= threading.Thread(target=self.start_aggregation, args=())
                            th_start_aggregate.start()
                        else:
                            self.command_send((self.up_server,self.port),receive_data.decode('utf-8'))#不是根节点就向上传递请求
                    elif(json_data['order']=="aggregation_confirm" and self.state!=2):
                        th_start_aggregate= threading.Thread(target=self.start_aggregation, args=())
                        th_start_aggregate.start()
                    elif (json_data['order']=="local_trained_ready"):#下层准备发送
                        self.prepare_receive_model(json_data)
                    elif (json_data['order']=="update_model_ready"):#上层准备更新
                        self.prepare_receive_new_model(json_data)
                        start_update=threading.Thread(target=self.start_update, args=())
                        start_update.start()
                    elif (json_data['order']=="start_send_data"):#开始发送
                        self.flag_start_send=True
                    elif(json_data['order']=="aggregation_resend"):
                        self.resend_local_model(json_data['rsend_id_list'],"aggregation_data",client[0])
                    elif(json_data['order']=="update_resend"):
                        self.resend_local_model(json_data['rsend_id_list'],"update_data",client[0])
                    elif(json_data['order']=="test_local_model"):#外部命令，测试模型准确度
                        self.test_local_model()
                    elif(json_data['order']=="read_model"):#外部命令，读取模型
                        self.local_model_read()
                    elif(json_data['order']=="reset_model"):#外部命令，重置模型
                        self.local_model_reset()
                    elif(json_data['order']=="save_model"):#外部命令，保存模型
                        self.local_model_save()
                    elif(json_data['order']=="send_aggregation_request"):#外部命令,向上请求聚合
                        self.udp_send((self.up_server,self.port),'{"order":"aggregation_request"}'.encode('utf-8'))
                    elif(json_data['order']=="change_topology"):#外部命令,修改拓扑结构
                        self.change_topology(json_data)
                    elif(json_data['order']=="update_part1"):#外部命令,更新推理模型
                        th_update_part1=threading.Thread(target=self.update_part1, args=("localhost",))
                        th_update_part1.start()
            except socket.timeout:
                pass
        print("stop listen")
        return
    def start(self):
        self.state_lock.acquire()
        self.state=1
        self.state_lock.release()
        self.running=True
        if not self.isdebug:
            pass
        else:
            print("It is debug mode")
        self.th_listen= threading.Thread(target=self.listen, args=())
        self.th_listen.start()
        print('Federated Learning v1.3.1 started at '+self.local_ip+":"+str(self.port))
        self.th_show= threading.Thread(target=self.show_state, args=())
        self.th_show.start()
        return
    def stop(self):
        self.running=False
        self.state_lock.acquire()
        self.state=0
        self.state_lock.release()
        self.th_listen.join()
        self.th_show.join()
        return
    def show_state(self):
        last_state=0
        while self.running:
            time.sleep(1)
            if self.state!=last_state:
                #print("state changes to:",self.state)
                last_state=self.state
            order_json=json.dumps({"order":"send_state","ip":s.local_ip,"isroot":self.isroot,"state":self.state,"up_server":self.up_server,"down_server":self.down_server})
            self.udp_send((self.cloud_server,self.port),order_json.encode('utf-8'))   
        return#向云服务器上传状态
    def change_topology(self,json_data):
        print("start to change_topology")
        if self.state!=1:
            text="ip:"+self.local_ip+",isn't in state 1,refuse to change topology"
            sendjson=json.dumps({"order":"print","text":""})
            self.udp_send((self.cloud_server,self.port),sendjson.encode('utf-8'))
            return
        self.state_lock.acquire()
        self.up_server=json_data['up_server']
        self.down_server=json_data['down_server']
        if self.up_server=="":
            self.isroot=True
        else:
            self.isroot=False
        self.clear_buffer()
        self.state_lock.release()
        sendjson=json.dumps({"order":"feedback","ip":self.local_ip ,"detail":"topology_changed"})
        self.udp_send((self.cloud_server,self.port),sendjson.encode('utf-8'))
        print("------------change_topology success----------")
        print("is_root:{}\nup_server:{}\ndown_server:{}".format(self.isroot,self.up_server,self.down_server))
        return

    def __init__(self,config_file):
        #获取计算机名称
        hostname=socket.gethostname()
        #获取本机IP
        self.local_ip=socket.gethostbyname(hostname)
        with open(config_file,'r',encoding='utf-8')as fp:
            json_data = json.load(fp)
            print('config file:',json_data)
            self.up_server=json_data['up_server']
            self.down_server=json_data['down_server']
            self.port=json_data['port']
            self.send_bytes_once=json_data['send_bytes_once']
            self.n_down_server=len(self.down_server)
            self.local_model=self.create_model()
            self.temp_model=self.create_model()
            if "local_ip" in json_data:
                self.local_ip=json_data["local_ip"]
            if "wait_time" in json_data:
                self.wait_time=json_data["wait_time"]
            if "isparticipant" in json_data:
                self.isparticipant=json_data["isparticipant"]
            if "aggregate_waite_time" in json_data:
                self.aggregate_waite_time=json_data["aggregate_waite_time"]
            if "send_speed" in json_data:
                self.send_speed=json_data["send_speed"]
            if "cloud_server" in json_data:
                self.cloud_server=json_data["cloud_server"]
            if "camera_path" in json_data:
                self.camera_path=json_data["camera_path"]
            if "mqtt_server" in json_data:
                self.mqtt_server=json_data["mqtt_server"]
            if "isdebug" in json_data:
                self.isdebug=json_data["isdebug"]
            if "send_image_channel" in json_data:
                self.send_image_channel=json_data["send_image_channel"]
            if json_data['local_model']==1:
                self.local_model_read()
            if self.up_server=="":
                self.isroot=True
        self.clear_buffer()
        return
    #---------------模型训练、预测、更新、存储、读取-----------------
    def local_model_train0(self):#本地训练
        print("start local_model_train")
        mnist = tf.keras.datasets.mnist
        (x_train, y_train),(x_test, y_test) = mnist.load_data()
        x_train, x_test = x_train / 255.0, x_test / 255.0
        data_choose=np.random.randint(0,6)
        self.local_model.fit(x_train, y_train, epochs=1, batch_size=10,verbose=1)
        #print(self.local_model.layers[1].get_weights()[0][0][0])
        print("local_model_train complete")
        return
    def local_model_train(self):#手势识别本地训练
        print("start local_model_train")
        image_list,lable_list=self.get_data_sets("./train_image")
        print("load data")
        x_train=self.read_data_sets(image_list)
        print("load data complete")
        y_train=tf.one_hot(lable_list,depth=11,dtype=tf.uint8)
        self.local_model.fit(x_train,y_train, epochs=1, batch_size=30,verbose=1)
        print("local_model_train complete")
        return
    def local_model_predict(self,addr):#传入图片地址，预测结果
        image = tf.io.read_file(addr)
        image = tf.image.decode_jpeg(image, channels=3)
        image = tf.image.resize(image, (50, 50))
        image = image/255.0
        x=np.zeros((1,50,50,3))
        x[0]=image
        ans=self.local_model.predict(x)#predict输出的是最后一层
        softmax=tf.keras.layers.Softmax()
        ans = tf.cast(ans, dtype=tf.float32)
        ans=softmax(ans)
        ans_index=np.argmax(ans)
        if ans.numpy()[0][ans_index]<0.5:
            print("no result")
            return -1
        print("predict result",ans_index)
        return ans_index
    def local_model_update(self):
        self.local_model.load_weights(".//model//UpdateModel.h5")
        return
    def local_model_read(self):
        self.local_model.load_weights('.//model//gestureModel.h5')
        print("model read")
        return
    def local_model_save(self):
        self.local_model.save_weights('.//model//gestureModel.h5')
        print("model saved")
        return
    def local_model_reset(self):
        self.local_model=self.create_model()
        print("model reset")
        return
    def test_local_model0(self):
        mnist = tf.keras.datasets.mnist
        (x_train, y_train),(x_test, y_test) = mnist.load_data()
        x_train, x_test = x_train / 255.0, x_test / 255.0
        print("*********test local model:************")
        print(self.local_model.evaluate(x_test, y_test)[1])
        #print(self.local_model.layers[1].get_weights()[0][0][0])
        return#mnist手写体识别
    def test_local_model(self):
        print("start local_model_train")
        image_list,lable_list=self.get_data_sets("./test_image")
        x_test=self.read_data_sets(image_list)
        y_test=tf.one_hot(lable_list,depth=11,dtype=tf.uint8)
        print(self.local_model.evaluate(x_test,y_test,batch_size=1)[1])
        return#手势识别
    def create_model0(self):#debug用的手写体识别的cnn
        model = tf.keras.models.Sequential([
        tf.keras.layers.Reshape((28,28,1),input_shape=(28, 28)),
        tf.keras.layers.Convolution2D(32, (5, 5), activation=tf.nn.relu),#    tf.keras.layers.Convolution2D(32, (5, 5), activation=tf.nn.relu),
        tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
        tf.keras.layers.Convolution2D(64, (5, 5), activation=tf.nn.relu),
        tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(512, activation=tf.nn.relu),
        tf.keras.layers.Dense(10, activation=tf.nn.softmax)
        ])

        model.compile(optimizer='SGD',
                      loss=tf.keras.losses.MeanSquaredError,
                      metrics=['accuracy'])
        self.model_data_layer_index=[1,3,6,7]
        self.n_data_layer=4
        return model
    def create_model(self):#手势识别的cnn
        model = tf.keras.models.Sequential([
            tf.keras.layers.Conv2D(16, 3, padding='same', activation='relu'),
            tf.keras.layers.MaxPooling2D(),
            tf.keras.layers.Conv2D(32, 3, padding='same', activation='relu'),
            tf.keras.layers.MaxPooling2D(),
            tf.keras.layers.Conv2D(64, 3, padding='same', activation='relu'),
            tf.keras.layers.MaxPooling2D(),
            tf.keras.layers.Flatten(),
            tf.keras.layers.Dense(128, activation='relu'),
            tf.keras.layers.Dense(11,activation=tf.nn.softmax)])
        model.build(input_shape=(None, 50, 50, 3))
        model.compile(optimizer='Adam',
                        loss=tf.keras.losses.MSE,
                        metrics=['accuracy'])
        model_depth=len(model.layers)
        model_data_layer_index=[]
        n_data_layer=0
        for i in range(0,model_depth):
            if model.layers[i].get_weights()!=[]:
                n_data_layer+=1
                model_data_layer_index.append(i)
        self.model_data_layer_index=model_data_layer_index
        self.n_data_layer=n_data_layer
        return model
    def get_layers_from_h5(self,path):
        self.temp_model.load_weights(path)
        temp_layers=[]
        for i in self.model_data_layer_index:
            temp_layers.append(self.temp_model.layers[i].get_weights())
        return temp_layers

    #---------------联邦学习相关-----------------
    def average_aggregate_model_root(self):
        all_weight=self.add_aggregate_model()#计算平均
        aggregate_layer=[]
        new_layer=[]
        for i in self.model_data_layer_index:
            aggregate_layer.append(self.temp_model.layers[i].get_weights())
        for i in range(0,self.n_data_layer):
            new_layer_data_i=[]
            for j in range(0,len(aggregate_layer[i])):
                new_layer_data_i.append(np.zeros(aggregate_layer[i][j].shape))
            new_layer.append(np.array(new_layer_data_i))
            new_layer[i]+=aggregate_layer[i]
            new_layer[i]=new_layer[i]/all_weight
            self.temp_model.layers[self.model_data_layer_index[i]].set_weights(new_layer[i])
        self.temp_model.save_weights('.//model//UpdateModel.h5')
        print("average_aggregate_model_root complete")
        return
    def add_aggregate_model(self):#计算平均模型，返回总模型权重
        #本地模型
        print("\nstart aggregation,this will take several minutes...")
        weight_all=0
        local_layer=[]
        new_layer=[]

        for i in self.model_data_layer_index:
            local_layer.append(self.local_model.layers[i].get_weights())


        for i in range(0,self.n_data_layer):#初始化new_layer
            new_layer_data_i=[]
            for j in range(0,len(local_layer[i])):
                new_layer_data_i.append(np.zeros(local_layer[i][j].shape))
            new_layer.append(np.array(new_layer_data_i))

        for i in range(0,self.n_down_server):
            if self.is_buffer_initialized[i]==1:
                path=".//model//aggregate_model_"+self.down_server[i]+".h5"
                r_layer_data=self.get_layers_from_h5(path)
                for j in range(0,self.n_data_layer):
                    for ii in range(0,len(local_layer[j])):
                        r_layer_data[j][ii]=np.array(r_layer_data[j][ii])
                    new_layer[j]+=np.array(r_layer_data[j])
                weight_all+=self.received_weight[i]

        if self.isparticipant:
            for j in range(0,self.n_data_layer):
                new_layer[j]+=local_layer[j]*self.local_weight#暂时默认本地模型权重为1
            weight_all+=self.local_weight
        for j in range(0,self.n_data_layer):
            self.temp_model.layers[self.model_data_layer_index[j]].set_weights(new_layer[j])
        self.temp_model.save_weights('.//model//AggregateModel.h5')#这里保存的梯度数据是不能直接用的，因为还没有除以权值
        print("\nadd_aggregate_model complete")
        return weight_all
    def update(self):#收齐数据
        received_timmer=time.time()
        last_receive_time=0
        last_receive_n=0
        while(self.state==4 and self.running):
            is_update_completed=True
            time.sleep(2)#每1s检查是否收齐，不加sleep会占用listen线程的资源，会很慢
            if last_receive_time==0:
                last_receive_time=time.time()#初始化，第一次收到数据才会执行start_update，所以此时已经有数据来了
            N=len(self.received_buffer_new)
            n=0
            for j in range(0,N):
                if self.received_buffer_new[j]!=[]:
                    n+=1
            if n<N:
                is_update_completed=False
                if n !=last_receive_n:
                    last_receive_n=n
                    last_receive_time=time.time()
                else:
                    if time.time()-last_receive_time>self.wait_time and n!=0:
                        resend_id_list=[]
                        for j in range(0,N):
                            if self.received_buffer_new[j]==[]:
                                resend_id_list.append(j)
                        resend_id_n=len(resend_id_list)#丢包的个数
                        print("packet lost",resend_id_n,",try to resend")
                        for ii in range(0,ceil(resend_id_n/50)):
                            resend_order_json=json.dumps({"order":"update_resend","rsend_id_list":resend_id_list[50*ii:50*(ii+1)]})
                            self.udp_send((self.up_server,self.port),resend_order_json.encode('utf-8'))
                            time.sleep(0.023)
                        resend_order_json=json.dumps({"order":"update_resend","rsend_id_list":resend_id_list[50*ceil(resend_id_n/50):]})
                        self.udp_send((self.up_server,self.port),resend_order_json.encode('utf-8'))
                        last_receive_time=time.time()
            if(is_update_completed):
                strb=b"".join(self.received_buffer_new)
                with open(".//model//UpdateModel.h5","wb") as f2:
                    f2.write(strb)
                del self.received_buffer_new
                self.received_buffer_new=None
                print("receive complete")
                break
        return#检查收齐数据
    def start_update(self):
        if not self.isroot:#非根节点需要接收更新数据
            th_update=threading.Thread(target=self.update, args=())
            th_update.start()
            sendjson=json.dumps({"order":"start_send_data"})#与接收下层聚合模型不同(那个在aggregate里确认)，这个直接开始传输就行
            self.command_send((self.up_server,self.port),sendjson)
            th_update.join()

        self.local_model_update()
        for i in range(0,self.n_down_server):#向下层发送更新数据
            print("ready to send new model to "+self.down_server[i])
            self.prepare_send_model((self.down_server[i],self.port),'.//model//UpdateModel.h5',"update_model_ready",1)
            while not self.flag_start_send:
                time.sleep(1)
            self.flag_start_send=False
            print("start send new model")
            self.send_model((self.down_server[i],self.port),"update_data")
            print("send complete")
        self.state_lock.acquire()
        self.state=1
        self.state_lock.release()
        return
    def aggregate(self):
        received_timmer=time.time()
        last_receive_time=[]
        last_receive_n=[]
        receive_complete=[]#0未开始 1传输中 2已完成
        for i in self.down_server:
            last_receive_time.append(0)
            last_receive_n.append(0)
            receive_complete.append(0)
        while(self.state==2 and self.running):
            is_aggregate_completed=True
            time.sleep(1)#每0.5s检查是否收齐，不加sleep会占用listen线程的资源，会很慢
            for i in range(0,self.n_down_server):#发送开始传输的命令
                if self.is_buffer_initialized[i]==1 and receive_complete[i]==0 and (1 not in receive_complete):
                    sendjson=json.dumps({"order":"start_send_data"})
                    self.command_send((self.down_server[i],self.port),sendjson)
                    receive_complete[i]=1
                if self.is_buffer_initialized[i]==0 or (self.is_buffer_initialized[i]==1 and receive_complete[i]!=2):
                    is_aggregate_completed=False

            for i in range(0,self.n_down_server):#i为服务器索引，j为buffer长度
                if self.is_buffer_initialized[i]==1 and receive_complete[i]==1:#只检查已经初始化（确定参与,且已经开始传输数据，而且没传输完）的缓存 
                    if last_receive_time[i]==0:
                        last_receive_time[i]=time.time()#初始化
                    N=len(self.received_buffer[i])
                    n=0
                    for j in range(0,N):
                        if self.received_buffer[i][j]!=[]:
                            n+=1
                    if n<N:
                        if n !=last_receive_n[i]:
                            last_receive_n[i]=n
                            last_receive_time[i]=time.time()
                        else:#如果一段时间没有新数据则重传
                            if time.time()-last_receive_time[i]>self.wait_time and n!=0:
                                resend_id_list=[]
                                for j in range(0,N):
                                    if self.received_buffer[i][j]==[]:
                                        resend_id_list.append(j)
                                resend_id_n=len(resend_id_list)#丢包的个数
                                print("packet lost",resend_id_n,",try to resend")
                                for ii in range(0,ceil(resend_id_n/50)):
                                    resend_order_json=json.dumps({"order":"aggregation_resend","rsend_id_list":resend_id_list[50*ii:50*(ii+1)]})
                                    self.udp_send((self.down_server[i],self.port),resend_order_json.encode('utf-8'))
                                    time.sleep(0.023)
                                resend_order_json=json.dumps({"order":"aggregation_resend","rsend_id_list":resend_id_list[50*ceil(resend_id_n/50):]})
                                self.udp_send((self.down_server[i],self.port),resend_order_json.encode('utf-8'))
                                last_receive_time[i]=time.time() 
                    else:
                        receive_complete[i]=2
                elif self.is_buffer_initialized[i]==0:#有子ip没有响应
                    if time.time()-received_timmer>self.aggregate_waite_time:#等待时间超过限制
                        print("waite"+self.down_server[i]+"overtime, give up this down server")
                        self.is_buffer_initialized[i]=-1


            if(is_aggregate_completed):
                for i in range(0,self.n_down_server):#保存模型
                    if self.is_buffer_initialized[i]==1:
                        strb=b"".join(self.received_buffer[i])
                        with open(".//model//aggregate_model_"+self.down_server[i]+".h5","wb") as f2:
                            f2.write(strb)
                del self.received_buffer
                self.received_buffer=None
                print("receive complete:",end='')
                self.state_lock.acquire()
                self.state=3
                self.state_lock.release()
        return#检查收齐数据
    def start_aggregation(self):#向下传递开始聚合的命令并进入聚合状态
        print("start_aggregation")
        self.state_lock.acquire()
        self.state=2
        self.state_lock.release()
        self.clear_buffer()
        if(self.n_down_server!=0):#有子节点才需要接收
            th_aggregate= threading.Thread(target=self.aggregate, args=())
            th_aggregate.start()
            for ds in self.down_server:
                self.command_send((ds,self.port),'{"order":"aggregation_confirm"}')
        else:#没有下层ip
            pass
        if(self.isparticipant):
            self.local_model_train()#开始本地训练
        if(self.n_down_server!=0):#有子节点才需要接收
            th_aggregate.join()
        #计算平均、向上传输
        self.state_lock.acquire()
        self.state=3
        self.state_lock.release()
        aggregate_empty=True#如果自己与子节点都不参与则不参与聚合
        if self.isparticipant:
            aggregate_empty=False
        else:
            for i in range(0,self.n_down_server):
                if self.is_buffer_initialized[i]==1:
                    aggregate_empty=False#自己不参与,但有子节点参与
        if not self.isroot:#非根节点
            if aggregate_empty:#本节点与子节点都不参与聚合
                resend_order_json=json.dumps({"order":"local_trained_ready","ip":s.local_ip,"N":0})
                self.command_send((self.up_server,self.port),resend_order_json)
            else:
                all_weight=self.add_aggregate_model()
                print("ready to send local model")
                self.prepare_send_model((self.up_server,self.port),'.//model//AggregateModel.h5',"local_trained_ready",all_weight)
                while not self.flag_start_send:
                    print("\rwaiting...",end="")
                    time.sleep(1)
                self.flag_start_send=False
                print("start send local model")
                self.send_model((self.up_server,self.port),"aggregation_data")#向上层发送聚合结果
                print("send local model complete")
                self.state_lock.acquire()#非root进入状态1等待分发新模型
                self.state=1
                self.state_lock.release()
        else:#是根节点
            if aggregate_empty:#本节点与子节点都不参与聚合
                print("this server and it's down server don't take part in the aggregration")
                self.state_lock.acquire()
                self.state=1
                self.state_lock.release()
                return
            self.average_aggregate_model_root()#计算平均
            #开始分发模型
            self.start_update()
        return

    #---------------传输数据-----------------
    def resend_local_model(self,rsend_id_list,order,ip):
        #已经调用过send_model并给self.data_bytes赋值过了，resend_local_model发送上一次send_model所发送的数据
        data_length=len(self.data_bytes)
        send_turns=ceil(data_length/self.send_bytes_once)
        sleep_time=50*self.send_bytes_once/self.send_speed*10e-6#控制发送速度,*50是发送50次数据等待一段时间
        for i in rsend_id_list:
            if i%50==1:
                time.sleep(sleep_time-0.0023)#0.0023是发送数据本身需要的时间，最大速度21.7Mb/s
            sendjson='{"order":"'+order+'","ip":"'+self.local_ip+'","N":'+str(send_turns)+',"data_id":'+str(i)+'}'
            self.udp_send((ip,self.port),sendjson.encode('utf-8')+self.data_bytes[i*self.send_bytes_once:(i+1)*self.send_bytes_once])
        return
    def send_model(self,addr,order):
        data_length=len(self.data_bytes)
        send_turns=ceil(data_length/self.send_bytes_once)
        sleep_time=50*self.send_bytes_once/self.send_speed*10e-6#控制发送速度,*50是发送50次数据等待一段时间
        for i in range(0,send_turns):#发送/接收过程
            if i%50==1:
                print('\r','{:.2f}'.format(i/send_turns),end='')
                time.sleep(sleep_time-0.0023)#0.0023是发送数据本身需要的时间，最大速度21.7Mb/s
            sendjson='{"order":"'+order+'","ip":"'+self.local_ip+'","data_id":'+str(i)+'}'
            self.udp_send(addr,sendjson.encode('utf-8')+self.data_bytes[i*self.send_bytes_once:(i+1)*self.send_bytes_once])
        #发送最后一个数据
        sendjson='{"order":"'+order+'","ip":"'+self.local_ip+'","data_id":'+str(send_turns-1)+'}'
        self.udp_send(addr,sendjson.encode('utf-8')+self.data_bytes[(send_turns-1)*self.send_bytes_once:])
        return
    def receive_model(self,json_data,receive_data2):#接收下层的聚合数据
        down_ip=json_data["ip"]
        if down_ip not in self.down_server:
            raise ValueError("receive data from"+down_ip+" ,but this ip is not my down server")
        ip_index=self.down_server.index(down_ip)
        self.received_buffer[ip_index][json_data["data_id"]]=receive_data2
        return
    def receive_update_model(self,json_data,receive_data2):#与receive_model不同，因为只需要接收上层的数据,且权重必为1
        self.received_buffer_new[json_data["data_id"]]=receive_data2
        return
    def prepare_receive_model(self,json_data):#该节点已经准备好传输，初始化缓存
        down_ip=json_data["ip"]
        N=int(json_data["N"])
        if down_ip not in self.down_server:
            raise ValueError("receive data from"+down_ip+" ,but this ip is not my down server")
        ip_index=self.down_server.index(down_ip)
        if N==0:#表示该IP不参与聚合
            self.received_weight[ip_index]=0
            self.is_buffer_initialized[ip_index]=-1
            print(down_ip+" don't take part in the aggregation")
        else:#参与聚合则初始化缓存
            for i in range(0,N):#初始化该下层ip的缓存
                self.received_buffer[ip_index].append([])
            self.is_buffer_initialized[ip_index]=1
            if "weight" not in json_data:
                raise RuntimeWarning(down_ip+"dosn't have weight data,defoult to 1")
                self.received_weight[ip_index]=1
            else:
                self.received_weight[ip_index]=int(json_data["weight"])
        return
    def prepare_receive_new_model(self,json_data):
        up_ip=json_data["ip"]
        N=int(json_data["N"])
        if up_ip != self.up_server:
            raise ValueError("receive new model data from"+up_ip+"but it is not my up server")
        if self.state ==1:#只在状态1接收更新
            self.state_lock.acquire()
            self.state=4
            self.state_lock.release()
            self.clear_buffer()
            for i in range(0,N):#初始化缓存
                self.received_buffer_new.append([])
        else:
            raise RuntimeWarning("refuse to receive new model while I am not in state 1")
        return
    def prepare_send_model(self,addr,path,order,all_weight):#准备传输模型
        with open(path,"rb") as f:
            self.data_bytes = f.read()
        data_length=len(self.data_bytes)
        send_turns=ceil(data_length/self.send_bytes_once)

        send_dic={"order":order,"weight":str(all_weight),"N":str(send_turns),"ip":self.local_ip}
        sned_json=json.dumps(send_dic)
        ans=self.command_send(addr,sned_json)
        return
    #---------------基本工具-----------------
    def clear_buffer(self):#初始化缓存
        self.received_buffer_new=[]

        self.received_buffer=[]#接收模型的缓存,是list（下层服务器个数，数据包发送总数）
        self.is_buffer_initialized=[]#0未初始化不知道是否参与，1已经初始化，-1不参与
        self.received_weight=[]#接收各个ip下数据的权重，在发送最后一个数据时发送
        for i in self.down_server:
            self.received_buffer.append([])
            self.is_buffer_initialized.append(0)
            self.received_weight.append(0)
        self.n_down_server=len(self.down_server)
        return
    def udp_send(self,address,message):
        client_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        client_socket.sendto(message, address) #将msg内容发送给指定接收方
        client_socket.close();
        return
    def command_send(self,address,message):#需要回执的命令 成功则返回true
        message_dic=json.loads(message)
        random_id=np.random.randint(1e6)
        message_dic['feedback_id']=random_id
        message_dic['port']=self.port
        message=json.dumps(message_dic)
        self.udp_send(address,message.encode("utf-8"))
        start_time=time.time()
        check_time=0.1#检查是否收到回执的间隙
        resend_time=0.5#重发消息的间隙
        resend_n=round(resend_time/check_time)
        n=0
        while(time.time()-start_time<3):
            time.sleep(check_time)
            n+=1
            if random_id in self.command_id:
                return True
            elif n==resend_n:
                self.udp_send(address,message.encode("utf-8"))#没有相应，重发命令
                n=0
        print("node {} is unresponsive, task id:{}".format(address,random_id))
        return False
    def feedback(self,address,json_data): #已经处理过的命令返回False，新命令返回True
        id=json_data['feedback_id']
        port=json_data['port']
        if "order" in json_data:#是feedback发送来的，不是新命令
            if json_data["order"]=="feedback":
                self.command_id.add(id)
                return False#回执不需要处理
        #可能是新命令
        send_dic={"order":"feedback","feedback_id":id,"port":port}
        sned_json=json.dumps(send_dic)

        self.udp_send((address[0],port),sned_json.encode('utf-8'))
        if id in self.command_id:
            return False
        self.command_id.add(id)
        return True
    #---------------应用相关-----------------
    def get_data_sets(self,addr):
        # 存放图片类别和标签的列表：第0类
        list_0 = []
        label_0 = []
        # 存放图片类别和标签的列表：第1类
        list_1 = []
        label_1 = []
        # 存放图片类别和标签的列表：第2类
        list_2 = []
        label_2 = []
        # 存放图片类别和标签的列表：第3类
        list_3 = []
        label_3 = []
        # 存放图片类别和标签的列表：第4类
        list_4 = []
        label_4 = []
        # 存放图片类别和标签的列表：第5类
        list_5 = []
        label_5 = []
        # 存放图片类别和标签的列表：第6类
        list_6 = []
        label_6 = []
        # 存放图片类别和标签的列表：第6类
        list_7 = []
        label_7 = []
        # 存放图片类别和标签的列表：第8类
        list_8 = []
        label_8 = []
        # 存放图片类别和标签的列表：第9类
        list_9 = []
        label_9 = []
        # 存放图片类别和标签的列表：第10类
        list_10 = []
        label_10 = []
        for file in os.listdir(addr):  # 获得file_dir路径下的全部文件名
            image_file_path = addr+"/"+file
            for image_name in os.listdir(image_file_path):
                image_name_path = os.path.join(image_file_path, image_name)
                # 将图片存放入对应的列表

                if image_file_path[-2:] == '10':
                    list_10.append(image_name_path)
                    label_10.append(10)
                elif image_file_path[-1:] == '0':
                    list_0.append(image_name_path)
                    label_0.append(0)
                elif image_file_path[-1:] == '1':
                    list_1.append(image_name_path)
                    label_1.append(1)
                elif image_file_path[-1:] == '2':
                    list_2.append(image_name_path)
                    label_2.append(2)
                elif image_file_path[-1:] == '3':
                    list_3.append(image_name_path)
                    label_3.append(3)
                elif image_file_path[-1:] == '4':
                    list_4.append(image_name_path)
                    label_4.append(4)
                elif image_file_path[-1:] == '5':
                    list_5.append(image_name_path)
                    label_5.append(5)
                elif image_file_path[-1:] == '6':
                    list_6.append(image_name_path)
                    label_6.append(6)
                elif image_file_path[-1:] == '7':
                    list_7.append(image_name_path)
                    label_7.append(7)
                elif image_file_path[-1:] == '8':
                    list_8.append(image_name_path)
                    label_8.append(8)
                elif image_file_path[-1:] == '9':
                    list_9.append(image_name_path)
                    label_9.append(9)



        # 合并数据
        image_list = np.hstack((list_0, list_1, list_2, list_3, list_4, list_5, list_6, list_7, list_8, list_9 ,list_10))
        label_list = np.hstack((label_0, label_1, label_2, label_3, label_4, label_5, label_6, label_7, label_8, label_9,label_10))

        # 利用shuffle打乱数据
        temp = np.array([image_list, label_list])
        temp = temp.transpose()  # 转置
        np.random.shuffle(temp)
        temp = temp.transpose()  # 转置
        image_list=temp[0]
        label_list=temp[1]
        #label_list=label_list.astype(np.float)
        #label_list=label_list.astype(np.int)
        return image_list ,np.array(label_list,dtype=np.uint8)
    def read_data_sets(self,image_list):
        x_train=np.zeros((len(image_list),50,50,3))
        for i in range(0,len(image_list)):
            image = tf.io.read_file(image_list[i])
            image = tf.image.decode_jpeg(image, channels=3)
            image = tf.image.resize(image, (50, 50))
            image = image/255.0
            x_train[i]=image
        return x_train
    def update_part1(self,ip):
        print("ready to send new model to "+ip)
        self.local_model.save_weights('.//model//update_to_part1.h5')
        self.prepare_send_model((ip,self.port+1),'.//model//update_to_part1.h5',"update_model_ready",1)
        while not self.flag_start_send:
            time.sleep(1)
        self.flag_start_send=False
        print("start send new model")
        self.send_model((ip,self.port+1),"update_data")
        print("send complete")
        return

if __name__ == '__main__':
    s=FLserver("config/config.json")
    s.start()
    s.local_model_read()






